import backtrader as bt
import numpy as np

from math_util.math_tool import cal_slope


class Check_bottom(bt.Indicator):
    lines = ('is_bottom',)
    params = dict(lookback=252, percentage=0.1)
    def __init__(self):
        data_lowest = bt.indicators.Lowest(self.data.low, period=self.params.lookback)
        data_highest = bt.indicators.Highest(self.data.high, period=self.params.lookback)
        data_position = data_lowest + self.p.percentage * (data_highest - data_lowest)
        self.lines.is_bottom = self.data.close < data_position
        
class BBand_converge(bt.Indicator):
    lines = ('is_converge',)
    params = dict(period=20, devfactor=2, converge_num=5,only_decline=True,slope=None)
    
    def __init__(self):
        bband = bt.indicators.BollingerBands(self.data.close,
                period=self.p.period, devfactor=self.p.devfactor)
        self.bband=bband
        self.bband_width = bband.lines.top - bband.lines.bot
        if self.p.converge_num:
            self.addminperiod(self.p.period + self.p.converge_num)
        self.min_width = None
        self.state='none'
        self.is_converge=False
    def next(self):
        if self.state=='none':
            if all(self.bband_width[-i] < self.bband_width[-(i + 1)] for i in range(self.p.converge_num)) and \
                self.bband.lines.mid[0] < self.bband.lines.mid[-self.p.converge_num]:
                if self.p.slope is None:
                    self.is_converge = True
                    self.min_width=self.bband_width[0]
                    self.state = 'best'
                else:
                    slope = cal_slope((np.array(self.bband_width.get(size=self.p.converge_num)))/self.bband_width[0])
                    if slope < -abs(self.p.slope):
                        self.is_converge = True
                        self.min_width=self.bband_width[0]
                        self.state = 'best'
                
        elif self.state=='best':
            if self.bband_width[0] < self.min_width:
                self.min_width=self.bband_width[0]
            elif self.bband_width[0] < self.min_width*1.1:
                self.state='continue'
            else:
                self.is_converge=False
                self.state='none'
        elif self.state=='continue':
            if self.bband_width[0] > self.min_width*1.1:
                self.is_converge=False
                self.state='none'
            elif self.bband_width[0] < self.min_width:
                self.min_width=self.bband_width[0]
        self.lines.is_converge[0] = self.is_converge
        
            
class OBV(bt.Indicator):
    lines = ('obv',)

    def __init__(self):
        self.addminperiod(2)

    def next(self):
        # check if self.lines.obv[-1] is NaN
        if self.lines.obv[-1] != self.lines.obv[-1]:
            self.lines.obv[-1] = 0
        if self.data.close[0] > self.data.close[-1]:
            self.lines.obv[0] = self.lines.obv[-1] + self.data.volume[0]
        elif self.data.close[0] < self.data.close[-1]:
            self.lines.obv[0] = self.lines.obv[-1] - self.data.volume[0]
        else:
            self.lines.obv[0] = self.lines.obv[-1]


        
class VolBurst(bt.Indicator):
    params = dict(period=20, threshold=3)
    lines = ('rv','rv_signal')
    def __init__(self):
        self.lines.rv = self.data.volume / bt.indicators.SMA(self.data.volume, period=self.p.period)
        self.lines.rv_signal = self.lines.rv > self.p.threshold
        
class Result_ind(bt.Indicator):
    lines = ('signal',)
    params = dict(strategy=None)

    def __init__(self):
        self.strategy = self.p.strategy

    def next(self):
        condition = self.strategy.condition  # 根据策略的状态来计算条件
        self.lines.signal[0] = 1.0 if condition else 0.0